import torch
import torch.nn as nn
import torchvision.models as models


class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.conv1 = nn.Sequential(nn.Conv2d(1, 64, 6, stride=2, padding=4),
                                   nn.BatchNorm2d(64),
                                   nn.LeakyReLU())
        self.conv2 = nn.Sequential(nn.Conv2d(64, 64, 4, stride=2, padding=2),
                                   nn.BatchNorm2d(64),
                                   nn.LeakyReLU())
        self.conv3 = nn.Sequential(nn.Conv2d(64, 64, 4, stride=2, padding=2),
                                   nn.BatchNorm2d(64),
                                   nn.LeakyReLU())
        self.conv4 = nn.Sequential(nn.Conv2d(64, 64, 4, stride=2, padding=2),
                                   nn.BatchNorm2d(64),
                                   nn.LeakyReLU())
        self.conv5 = nn.Sequential(nn.Conv2d(64, 64, 4, stride=2, padding=2),
                                   nn.BatchNorm2d(64),
                                   nn.LeakyReLU())
        self.conv6 = nn.Sequential(nn.Conv2d(64, 64, 4, stride=2, padding=2),
                                   nn.BatchNorm2d(64),
                                   nn.LeakyReLU())
        self.conv7 = nn.Sequential(nn.Conv2d(64, 64, 4, stride=2, padding=2),
                                   nn.BatchNorm2d(64),
                                   nn.LeakyReLU())
        self.conv8 = nn.Sequential(nn.Conv2d(64, 64, 4, stride=1),
                                   # nn.BatchNorm2d(64),
                                   nn.LeakyReLU())
        self.deconv1 = nn.Sequential(nn.ConvTranspose2d(64, 64, 4, stride=1),
                                     # nn.BatchNorm2d(64),
                                     nn.ReLU(),
                                     nn.Dropout(0.5))
        self.deconv2 = nn.Sequential(nn.ConvTranspose2d(64, 64, 4, stride=2, padding=2),
                                     nn.BatchNorm2d(64),
                                     nn.ReLU(),
                                     nn.Dropout(0.5))
        self.deconv3 = nn.Sequential(nn.ConvTranspose2d(64, 64, 4, stride=2, padding=2),
                                     nn.BatchNorm2d(64),
                                     nn.ReLU(),
                                     nn.Dropout(0.5))
        self.deconv4 = nn.Sequential(nn.ConvTranspose2d(64, 64, 4, stride=2, padding=2),
                                     nn.BatchNorm2d(64),
                                     nn.ReLU())
        self.deconv5 = nn.Sequential(nn.ConvTranspose2d(64, 64, 4, stride=2, padding=2),
                                     nn.BatchNorm2d(64),
                                     nn.ReLU())
        self.deconv6 = nn.Sequential(nn.ConvTranspose2d(64, 64, 4, stride=2, padding=2),
                                     nn.BatchNorm2d(64),
                                     nn.ReLU())
        self.deconv7 = nn.Sequential(nn.ConvTranspose2d(64, 64, 4, stride=2, padding=2),
                                     nn.BatchNorm2d(64),
                                     nn.ReLU())
        self.deconv8 = nn.Sequential(nn.ConvTranspose2d(64, 1, 6, stride=2, padding=4),  # 论文中output是3
                                     nn.Tanh())

    def forward(self, input_img):
        x1 = self.conv1(input_img)
        x2 = self.conv2(x1)
        x3 = self.conv3(x2)
        x4 = self.conv4(x3)
        x5 = self.conv5(x4)
        x6 = self.conv6(x5)
        x7 = self.conv7(x6)
        x8 = self.conv8(x7)
        x9 = self.deconv1(x8)
        x10 = self.deconv2(x9+x7)
        x11 = self.deconv3(x10+x6)
        x12 = self.deconv4(x11+x5)
        x13 = self.deconv5(x12+x4)
        x14 = self.deconv6(x13+x3)
        x15 = self.deconv7(x14+x2)
        x = self.deconv8(x15+x1)
        return x


class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.dis = nn.Sequential(nn.Conv2d(2, 48, 4, stride=2),
                                 nn.BatchNorm2d(48),
                                 nn.LeakyReLU(),
                                 nn.Conv2d(48, 96, 4, stride=2),
                                 nn.BatchNorm2d(96),
                                 nn.LeakyReLU(),
                                 nn.Conv2d(96, 192, 4, stride=2),
                                 nn.BatchNorm2d(192),
                                 nn.LeakyReLU(),
                                 nn.Conv2d(192, 384, 4, stride=1),
                                 nn.BatchNorm2d(384),
                                 nn.LeakyReLU(),
                                 nn.Conv2d(384, 1, 4, stride=1),
                                 nn.Tanh())

    def forward(self, *input):
        x = torch.cat((input[0], input[1]), 1)
        return self.dis(x)


class Vgg(nn.Module):
    def __init__(self):
        super(Vgg, self).__init__()
        vgg16 = models.vgg16(pretrained=False)
        self.features = nn.Sequential(*list(vgg16.features.children())[0:9])

    def forward(self, input_img):
        x = torch.cat((input_img, input_img, input_img), 1)
        return self.features(x)
